# aem.matrix -- Matrix solver classes
#
# Vic Kelson, SFWMD

from numpy import *
#from LinearAlgebra import *		#maybe replace by "import numpy.linalg as linalg"

class MatrixRow:
	"""Matrix row object stores a row and right hand side
	
	Contains information for one row (or part of one row) of a full matrix.
	
	Attributes:
	
	  - A: the element(s)  A_{ij} of the i'th row of the matrix, stored as a list,
	    e.g. [ A_{i,1}, A_{i,2}, ... ]
	  - b: a scalar containing B_{i}, the i'th element of the known vector
	
	Methods:
	
	  - __init__(self,element,index=0,A=[],b=0.0):
	    Constructor.  Defaults to an empty a row vector and RHS=0
	  - concat(self,matRow):
	    Concatenates matRow onto 'self'.

	    Concatenation is defined as: -- matRowC = matRowA + matRowB, 
	    where -- matRowC.a = matRowA.a + matRowB.a
	    and -- matRowC.b = matRowA.b + matRowB.b

	    This is convenient for accumulation during the construction of a complete
	    matrix row; the getMatrixContribution() methods return a MatrixRow for the
	    single element, while getMatrixRows() methods sum getMatrixContribution()
	    results across the solution.
	"""
	def __init__(self,A=[],b=0.0):
		self.A = A
		self.b = b

class MatrixColumn:
	"""Matrix colum object stores reference to column
	
	Contains information for one column (or part of one column) of a matrix.
	
	Attributes:
	
	  - element: the element which corresponds to the matrix row
	  - index: the parameter index into the element for the row
	
	Methods:
	
	  - __init__(self,element,index=0): Constructor.  Defaults to index=0
	"""
	def __init__(self,element,index=0):
		self.element=element
		self.index = index

class Solver:
	"""Solves system for a given list of MatrixRow objects
	
	Solves a system of equations, using a list of MatrixRow objects as the
	basis for the solution.
	
	Attributes:
	
	- A: the matrix
	- B: the known vector
	- X: the solution vector
	"""
	def __init__(self,rows):
		'''Constructor

		The map function is a nifty Python function that can be used
		to perform an operation (defined with the lambda function) over
		all components of a list. In this case all A components are taken from the
		MatrixRow object and put in a list. This list of lists of true matrix
		rows is then converted to an array, so that it can be solved using matrix
		solvers provided in numpy'''
		self.rows = rows
		self.A = array(map(lambda(z):z.A,rows),dtype='f')
		self.B = array(map(lambda(z):[z.b],rows),dtype='f')
	def setRHS(self,b):
		'''Distract right hand side from MatrixRow objects'''
		self.B = array(map(lambda(z):[z],b),dtype='f')
	def solve(self):
		'''Obtain a solution and store result
		
		Solves the system A * X = B. If the number of rows of A is larger than
		the number of columns, performs a least-squares solution.  Stores the new
		X array into the x member of self.rows[i]'''
		if len(self.rows) > 0 :   # added test to prevent error
		                          # when there is nothing to solve - wjz
			if self.A.shape[0] == self.A.shape[1]:
				# NOTE the solver throws an error "singular system"
				#      when the matrix is singular and there is no
				#      solution, but not when there are infinitely
				#      many solutions (in the latter case it just
				#      returns one solution without any warning)     - wjz
				return linalg.solve(self.A,self.B)
			else:
				return linalg.lstsq(self.A,self.B)[0]
		else :                    # return dummy when there
			return []              # is nothing to solve for - wjz
